package com.hzw.saas.web.sso.security;

import com.hzw.saas.common.security.pojo.dto.SaasUser;
import org.springframework.security.oauth2.common.DefaultOAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.TokenEnhancer;

import java.util.Map;

/**
 * @author sonam
 * @sine 2021/11/2 10:14 上午
 */
public class CustomAdditionalInformation implements TokenEnhancer {

    public static final String USER_ID = "userId";

    @Override
    public OAuth2AccessToken enhance(OAuth2AccessToken accessToken, OAuth2Authentication authentication) {
        Object principal = authentication.getPrincipal();
        if(principal instanceof SaasUser) {
            Map<String, Object> info = accessToken.getAdditionalInformation();
            info.put(USER_ID, ((SaasUser) principal).getUserId());
            ((DefaultOAuth2AccessToken) accessToken).setAdditionalInformation(info);
        }
        return accessToken;
    }
}
